# python3

import requests
import json
import arrow
import tushare as ts
from tqdm import tqdm
import time
import tempfile
import csv
import sys
sys.setrecursionlimit(10000) 
import logging
from functools import wraps
from collections import defaultdict
from config import get_header

def get_log():
    log = logging.getLogger(__name__)
    ft = logging.Formatter('%(asctime)s - %(name)s - %(funcName)s - %(levelname)s - %(message)s','%Y-%m-%d %H:%M:%S')

    fh = logging.FileHandler('%s.log'%__name__)
    fh.setFormatter(ft)
    fh.setLevel(logging.WARNING)

    sh = logging.StreamHandler()
    sh.setFormatter(ft)
    sh.setLevel(logging.DEBUG)

    log.addHandler(fh)
    log.addHandler(sh)

    log.setLevel(logging.DEBUG)
    return log

log = get_log()

class setLog:
    @classmethod
    def request(cls, retry=3):
        def decorate(func):
            @wraps(func)
            def wrapper(*args, **kwargs):
                exception_count = defaultdict(lambda:0)
                for _ in range(retry):
                    try:
                        return func(*args, **kwargs)
                    except requests.exceptions.RequestException as e:
                        exception_count['RequestException'] += 1
                        log.warning(e)
                        time.sleep(5)
                    except AssertionError as e:
                        exception_count['AssertionError'] += 1
                        log.warning('%s %s'%(str(e),str(args if args else kwargs)))
                        if exception_count['AssertionError'] >= 1:
                            log.error('%s %s time over'%(str(e),str(args if args else kwargs)))
                            return
                        time.sleep(5)
                    except TypeError as e:
                        exception_count['TypeError'] += 1
                        log.warning('%s %s'%(str(e),str(args if args else kwargs)))
                        if exception_count['TypeError'] >= 1:
                            log.error('%s %s TypeError'%(str(e),str(args if args else kwargs)))
                            return
                        time.sleep(1)
                    except Exception as e:
                        exception_count['Exception'] += 1
                        log.exception(e)
                        time.sleep(5)
                else:
                    log.error('request fail %s'%str((args if args else kwargs)))
            return wrapper
        return decorate

class stock:
    def parse_dic_data(self, dic, symbol, arrnow, arrupdate=None):
        dic['_id'] = symbol + '_' + arrnow.format('YYYYMMDD')
        dic['symbol'] = symbol
        dic['update'] = arrupdate.datetime if arrupdate else arrnow.datetime

    def __repr__(self):
        return self.__class__.__name__

class Neeq(stock):

    def __init__(self):
        self.totalPages = None
        self.org_count = None
        self.orgs = []
            
    def get_neeq_coms(self):
        def parse_org(dic):
            self.parse_dic_data(dic,dic.get('xxzqdm'),arrow.now().floor('hour'))
            return dic
        self._get_neeq_coms()

        self.orgs = list(map(parse_org, self.orgs))
        return self.orgs

    @setLog.request()
    def _get_neeq_coms(self,page=0):
        # http://www.neeq.com.cn/nq/listedcompany.html
        neep_total_url = 'http://www.neeq.com.cn/nqxxController/nqxx.do'
        req = requests.get(neep_total_url,params={'page':page,'typejb':'T'},timeout=10)
        assert req.status_code == 200 ,requests.exceptions.RequestException
        html = req.text.strip('null()')
        j_data = json.loads(html)[0]
        if not self.totalPages:
            self.totalPages = j_data.get('totalPages')
            self.org_count = j_data.get('totalElements')

        if len(j_data.get('content')) > 0:
            self.orgs.extend(j_data.get('content'))
        print(len(self.orgs), end='\r')

        if page<self.totalPages-1:
            page += 1
            self._get_neeq_coms(page)

    @setLog.request()
    def get_neeq_detail(self,symbol):
        # http://www.neeq.com.cn/nq/detailcompany.html?companyCode=430003&typeId=1&typename=G&zrlx=%E5%8D%8F%E8%AE%AE
        neep_detail_url = 'http://www.neeq.com.cn/nqhqController/detailCompany.do'#zqdm
        req = requests.get(neep_detail_url,params={'zqdm':symbol},timeout=10)
        assert req.status_code == 200 ,requests.exceptions.RequestException
        dic = json.loads(req.text.strip('null()'))
        self.parse_dic_data(dic,symbol,arrow.now().floor('hour'))
        return dic



class Snowball(stock):

    def __init__(self):
        self.session = requests.session()
        self.session.headers.update(get_header())
        self.session.get('https://xueqiu.com',timeout=10)

    @setLog.request()
    def get_basic_info(self,symbol,typ):
        '''
        https://xueqiu.com/S/OC430002/GSJJ
            :typ: oc is neeq
        '''
        url = 'https://xueqiu.com/stock/f10/compinfo.json'
        req = self.session.get(url,params={'symbol':'%s%s'%(typ,symbol)},timeout=10) 
        assert req.status_code == 200 ,requests.exceptions.RequestException
        assert '请输入验证码' not in req.text, '请输入验证码'
        dic = req.json().get('tqCompInfo')
        self.parse_dic_data(dic,symbol,arrow.now().floor('hour'))
        return dic

    @setLog.request()
    def get_capital_structure(self,symbol,typ):
        '''
        https://xueqiu.com/S/OC430002/GBJG
        typ: oc is neeq
        '''
        url = 'https://xueqiu.com/stock/f10/shareschg.json'
        req = self.session.get(url,params={'symbol':'%s%s'%(typ,symbol),'size':500},timeout=10)
        assert req.status_code == 200 ,requests.exceptions.RequestException
        assert '请输入验证码' not in req.text, '请输入验证码'
        lis = req.json().get('list')

        def parse_dic(dic):
            exrightdate = arrow.get(dic.get('exrightdate'),'YYYYMMDD').replace(tzinfo='local')
            self.parse_dic_data(dic,symbol,exrightdate,arrow.now().floor('hour'))
            dic['exrightdate'] = exrightdate.datetime
            return dic

        lis = list(map(parse_dic,lis)) 

        assert lis ,'ret lis is None'
        return lis

    @setLog.request()
    def get_financial_indicators(self,symbol,typ):
        '''
        https://xueqiu.com/S/OC430002/ZYCWZB
        typ: oc is neeq
        '''
        url = 'https://xueqiu.com/stock/f10/finmainindex.json'
        req = self.session.get(url,params={'symbol':'%s%s'%(typ,symbol),'size':500},timeout=10)
        assert req.status_code == 200 ,requests.exceptions.RequestException
        assert '请输入验证码' not in req.text, '请输入验证码'
        lis = req.json().get('list')

        def parse_dic(dic):
            begindate = arrow.get(dic.get('reportdate'),'YYYYMMDD').replace(tzinfo='local')
            self.parse_dic_data(dic,symbol,begindate, arrow.now().floor('hour'))
            dic['reportdate'] = begindate.datetime
            return dic

        lis = list(map(parse_dic,lis))

        assert lis,'ret lis is None'
        return lis

    @setLog.request()
    def get_h_stock(self,symbol,typ):
        '''
        https://xueqiu.com/S/OC430002
        typ: oc is neeq
        '''
        url = 'https://xueqiu.com/stock/forchartk/stocklist.json'
        req = self.session.get(url,params={'symbol':'%s%s'%(typ,symbol),'period':'1day','type':'normal'},timeout=10)
        assert req.status_code == 200 ,requests.exceptions.RequestException
        assert '请输入验证码' not in req.text, '请输入验证码'
        lis = req.json().get('chartlist')
        def parse_dic(dic):
            begindate = arrow.get(dic.get('time'),'MMM D HH:mm:ss Z YYYY')
            self.parse_dic_data(dic,symbol,begindate, arrow.now().floor('hour'))
            dic['time'] = begindate.datetime
            return dic

        lis = list(map(parse_dic,lis))
        assert lis ,'ret lis is None'
        return lis

class TencentStock(stock):

    def __init__(self):
        self.session = requests.session()
        self.session.headers.update(get_header())
    @setLog.request()
    def get_basic_info(self,symbol,typ):
        '''
        http://gu.qq.com/nq430002/jbzl
        :typ: nq is neep
        only neeq
        '''
        url = 'http://web.ifzq.gtimg.cn/stock/corp/Basicinfo/getBasicinfo'
        req = self.session.get(url,params={'symbol':'%s%s'%(typ,symbol)},timeout=10)
        assert req.status_code == 200 ,requests.exceptions.RequestException
        dic = req.json().get('data')
        self.parse_dic_data(dic,symbol,arrow.now().floor('hour'))
        return dic

class My_tushare(stock):

    def __init__(self):
        pass

    @setLog.request()
    def get_my_area_classified(self):
        def parse_dic(dic):
            self.parse_dic_data(dic,dic.get('code'), arrow.now().floor('hour'))
            return dic

        df = ts.get_area_classified()
        lis = json.loads(df.to_json(orient='records'))
        lis = list(map(parse_dic,lis))
        return lis


    @setLog.request()
    def get_my_h_data(self,symbol):
        def parse_dic(dic):
            dic_date = arrow.get(dic.get('date'),'YYYY-MM-DD').replace(tzinfo='local')
            self.parse_dic_data(dic, symbol,dic_date, arrow.now().floor('hour'))
            dic['date'] = dic_date.datetime
            return dic

        df = ts.get_h_data(symbol,start='2017-01-01',retry_count=10)
        fp = tempfile.TemporaryFile('w+t')
        df.to_csv(fp)
        fp.seek(0)
        reader = csv.DictReader(fp)
        lis = [row for row in reader]
        fp.close()
        lis = list(map(parse_dic,lis))
        return lis


if __name__ == '__main__':
    from pprint import pprint
    sb = Snowball()
    ret = sb.get_capital_structure('832513','OC')
    print(ret)
    # sb.get_basic_info(430002,'OC')
    # sb = TencentStock()
    # sb.get_basic_info(430002,'nq')
    # np = Neeq()   
    # mt = My_tushare()
    # ret = mt.get_my_h_data('600326')
    # pprint(ret)
    pass


